import numpy as np
import pandas as pd
import plotly.graph_objects as go


def left_gaussian_smoothing(data, sigma):
    smoothed = []
    for i in range(len(data)):
        kernel = np.exp(-np.square(np.arange(0, i + 1)[::-1]) / (2 * sigma**2))
        kernel /= kernel.sum()  # 归一化
        smoothed.append(np.dot(kernel, data[: i + 1]))
    return np.array(smoothed)


def draw():
    df = pd.read_parquet(
        "../cache/3m/ETH_USDT/1609430400000_1731081600000.parquet"
    )  # type: pd.DataFrame
    df = df.iloc[3000:5000]

    # b, a = butter(N=3, Wn=0.2, btype='low')
    # smoothed = filtfilt(b, a, df["close"])

    # smoothed = gaussian_filter1d(df["close"], sigma=1.4)
    smoothed = left_gaussian_smoothing(df["close"], sigma=2)
    # smoothed = talib.MA(df["close"], timeperiod=7)

    # smoothed = np.diff(df["close"], n=1)
    # smoothed = savgol_filter(df["close"], window_length=4, polyorder=2)
    # coeffs = pywt.wavedec(df["close"], 'db1', level=8)  # 小波分解
    # coeffs[-1] = np.zeros_like(coeffs[-1])  # 去除高频噪声
    # 清除更多高频系数，如 D1 和 D2
    # coeffs[-1] = np.zeros_like(coeffs[-1])  # 清除 D1
    # coeffs[-2] = np.zeros_like(coeffs[-2])  # 清除 D2
    # 阈值去噪：对每个细节系数设置阈值
    # threshold = 0.1 * max(abs(coeffs[-1]))  # 选择适当的阈值
    # coeffs = [pywt.threshold(c, threshold, mode='soft') if i != 0 else c for i, c in enumerate(coeffs)]
    # smoothed = pywt.waverec(coeffs, 'db1')  # 小波重构
    smoothed = pd.Series(smoothed, index=df.index)
    diff1 = smoothed - smoothed.shift(1)
    diffpos = (diff1.shift(-1) >= 0) & (diff1.shift(0) < 0)
    diffneg = (diff1.shift(-1) <= 0) & (diff1.shift(0) > 0)
    # ma1 = df["close"].rolling(window=7, min_periods=7, center=True).mean().shift(1)
    # ma2 = df["close"].rolling(window=14, min_periods=14, center=True).mean().shift(7)
    ma1 = left_gaussian_smoothing(df["close"], sigma=5)
    ma2 = left_gaussian_smoothing(df["close"], sigma=7)

    diff_choose = df[diffpos | diffneg]
    diff_change = (diff_choose["close"].shift(-1) - diff_choose["close"]) / (
        diff_choose["close"] + 1e-9
    )
    # diff_change = diff_change[np.abs(diff_change) > 0.005]

    diffpos[df.index.difference(diff_change.index)] = False
    diffneg[df.index.difference(diff_change.index)] = False
    df["peak_change"] = diff_change

    fig = go.Figure()
    fig.add_trace(
        go.Candlestick(
            x=df.index,  # 时间
            open=df["open"],  # 开盘价
            high=df["high"],  # 最高价
            low=df["low"],  # 最低价
            close=df["close"],  # 收盘价
        )
    )

    fig.add_trace(
        go.Scatter(
            x=df[diffpos].index,
            y=df[diffpos]["low"] - 1,
            mode="markers",  # 仅显示散点
            marker=dict(size=5, color="blue"),  # 散点大小和颜色
            name="pos",
        )
    )

    fig.add_trace(
        go.Scatter(
            x=df[diffneg].index,
            y=df[diffneg]["high"] + 1,
            mode="markers",  # 仅显示散点
            marker=dict(size=5, color="red"),  # 散点大小和颜色
            name="neg",
        )
    )

    fig.add_trace(
        go.Scatter(
            x=df.index,
            y=smoothed,
            name="拟合曲线",
            yaxis="y1",
            mode="lines",
        )
    )

    fig.add_trace(
        go.Scatter(
            x=df.index,
            y=ma1,
            name="ma1",
            yaxis="y1",
            mode="lines",
        )
    )

    fig.add_trace(
        go.Scatter(
            x=df.index,
            y=ma2,
            name="ma2",
            yaxis="y1",
            mode="lines",
        )
    )

    fig.update_layout(
        title="k线信号分析",
        xaxis_title="日期",
        yaxis_title="价格",
        xaxis_type="category",
        xaxis_rangeslider_visible=True,
        height=1000,
        yaxis2=dict(
            title="增长率",
            overlaying="y",
            side="right",
            fixedrange=False,
            zeroline=True,  # 显示x轴的0轴线
            zerolinewidth=2,  # 设置0轴线宽度
            zerolinecolor="gray",
        ),
        yaxis=dict(fixedrange=False),
    )
    fig.write_html("rub/price_analysis.html")


if __name__ == "__main__":
    draw()
